#include <cuda_runtime.h>
#include <stdio.h>

// CUDA 错误检查宏
#define CUDA_CHECK(call)                                                     \
    do                                                                       \
    {                                                                        \
        cudaError_t err = call;                                              \
        if (err != cudaSuccess)                                              \
        {                                                                    \
            fprintf(stderr, "CUDA error at %s %d: %s\n", __FILE__, __LINE__, \
                    cudaGetErrorString(err));                                \
            exit(EXIT_FAILURE);                                              \
        }                                                                    \
    } while (0)

// CUDA 内核函数
__global__ void conv2d_kernel(
    int IC, int IH, int IW,
    int OC, int OH, int OW,
    const float *__restrict__ inp,
    const float *__restrict__ weight,
    float *__restrict__ out)
{
    // 计算输出特征图的坐标
    int ow = blockIdx.x * blockDim.x + threadIdx.x;
    int oh = blockIdx.y * blockDim.y + threadIdx.y;
    int oc = blockIdx.z;

    if (ow >= OW || oh >= OH)
        return;

    float value = 0.0f;

    for (int ic = 0; ic < IC; ++ic)
    {
        for (int kh = 0; kh < 3; ++kh)
        {
            for (int kw = 0; kw < 3; ++kw)
            {
                int ih = oh + kh;
                int iw = ow + kw;
                float inp_val = inp[ic * IH * IW + ih * IW + iw];
                float weight_val = weight[oc * IC * 3 * 3 + ic * 3 * 3 + kh * 3 + kw];
                value += inp_val * weight_val;
            }
        }
    }

    out[oc * OH * OW + oh * OW + ow] = value;
}

// 主机函数，用于配置和启动CUDA内核
void conv2d_cuda(
    int IC, int IH, int IW,
    int OC,
    const float *d_inp,
    const float *d_weight,
    float *d_out)
{
    // 计算输出特征图的尺寸
    int OH = IH - 3 + 1; // padding=0, stride=1
    int OW = IW - 3 + 1;

    // 定义CUDA网格和块的维度
    dim3 blockDim(16, 16);
    dim3 gridDim((OW + blockDim.x - 1) / blockDim.x,
                 (OH + blockDim.y - 1) / blockDim.y,
                 OC); // 每个输出通道一个网格块

    // 启动CUDA内核
    conv2d_kernel<<<gridDim, blockDim>>>(
        IC, IH, IW,
        OC, OH, OW,
        d_inp,
        d_weight,
        d_out);

    // 检查内核是否有错误
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
}
